import ebooklib
from ebooklib import epub
from bs4 import BeautifulSoup
import os
import re
import html2text
import shutil
from urllib.parse import unquote

def clean_text(text):
    """清理文本，移除多余的空格和换行"""
    # 移除多余的空格
    text = re.sub(r'\s+', ' ', text)
    # 移除多余的空行
    text = re.sub(r'\n\s*\n', '\n\n', text)
    return text.strip()

def is_main_chapter(title):
    """判断是否是主章节"""
    # 匹配"第X章"或"第X部分"等格式
    patterns = [
        r'^第[一二三四五六七八九十百千万]+章',
        r'^第[一二三四五六七八九十百千万]+部分',
        r'^第\d+章',
        r'^第\d+部分'
    ]
    return any(re.match(pattern, title) for pattern in patterns)

def convert_img_to_markdown(html_content):
    """将 HTML 中的 img 标签转换为 Markdown 图片语法"""
    soup = BeautifulSoup(html_content, 'html.parser')
    for img in soup.find_all('img'):
        src = img.get('src', '')
        # 确保路径包含 images 目录
        if not src.startswith('images/'):
            src = f'images/{os.path.basename(src)}'
        alt = img.get('alt', '')
        title = img.get('title', '')
        # 创建 Markdown 图片语法
        markdown_img = f'![{alt}]({src} "{title}")'
        img.replace_with(markdown_img)
    return str(soup)

def get_chapter_content(soup):
    """从 BeautifulSoup 对象中提取完整的章节内容"""
    # 获取所有内容
    content = []
    for element in soup.find_all(['p', 'div', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6', 'pre', 'code', 'blockquote', 'img']):
        content.append(str(element))
    return '\n'.join(content)

def epub_to_markdown(epub_path, output_dir=None):
    """
    将 EPUB 文件转换为 Markdown 格式，每个章节保存为单独的文件
    
    Args:
        epub_path: EPUB 文件路径
        output_dir: 输出目录路径，如果为 None 则使用 EPUB 文件名创建目录
    """
    if not os.path.exists(epub_path):
        print(f"错误：文件 {epub_path} 不存在")
        return
    
    # 创建输出目录
    if output_dir is None:
        output_dir = os.path.splitext(epub_path)[0]
    os.makedirs(output_dir, exist_ok=True)
    
    # 创建图片目录
    images_dir = os.path.join(output_dir, 'images')
    os.makedirs(images_dir, exist_ok=True)
    
    try:
        # 读取 EPUB 文件
        book = epub.read_epub(epub_path)
        
        # 获取书籍标题
        title = book.get_metadata('DC', 'title')[0][0] if book.get_metadata('DC', 'title') else 'Untitled'
        
        # 创建转换器
        h = html2text.HTML2Text()
        h.ignore_links = False
        h.ignore_images = False
        h.body_width = 0  # 不限制行宽
        h.images_as_html = False  # 不将图片转换为 HTML 格式
        h.protect_links = True  # 保护链接不被转换
        h.wrap_links = False  # 不包装链接
        
        # 存储当前章节内容
        current_chapter = None
        current_content = []
        
        # 图片路径映射表
        image_map = {}
        
        # 处理所有项目
        for item in book.get_items():
            # 处理图片
            if item.get_type() == ebooklib.ITEM_IMAGE:
                # 获取图片文件名
                img_name = os.path.basename(unquote(item.get_name()))
                # 保存图片
                img_path = os.path.join(images_dir, img_name)
                with open(img_path, 'wb') as f:
                    f.write(item.get_content())
                # 存储映射关系，使用相对路径
                image_map[item.get_name()] = f'./images/{img_name}'
                print(f"已保存图片：{img_name}")
            
            # 处理文档
            elif item.get_type() == ebooklib.ITEM_DOCUMENT:
                # 获取章节内容
                content = item.get_content().decode('utf-8')
                
                # 使用 BeautifulSoup 解析 HTML
                soup = BeautifulSoup(content, 'html.parser')
                
                # 更新图片路径
                for img in soup.find_all('img'):
                    src = img.get('src', '')
                    if src in image_map:
                        img['src'] = image_map[src]
                        # 添加 alt 文本（如果有）
                        if not img.get('alt'):
                            img['alt'] = os.path.basename(src)
                
                # 获取所有标题
                headings = soup.find_all(['h1', 'h2', 'h3', 'h4', 'h5', 'h6'])
                
                for heading in headings:
                    heading_text = heading.get_text().strip()
                    
                    # 如果是主章节
                    if is_main_chapter(heading_text):
                        # 保存之前的章节（如果有）
                        if current_chapter:
                            # 创建安全的文件名
                            safe_filename = re.sub(r'[\\/*?:"<>|]', '_', current_chapter)
                            safe_filename = re.sub(r'\s+', '_', safe_filename)
                            
                            # 保存章节内容
                            output_path = os.path.join(output_dir, f"{safe_filename}.md")
                            with open(output_path, 'w', encoding='utf-8') as f:
                                f.write(f"# {current_chapter}\n\n")
                                # 使用 html2text 转换内容
                                html_content = '\n'.join(current_content)
                                # 将 img 标签转换为 Markdown 图片语法
                                html_content = convert_img_to_markdown(html_content)
                                markdown_content = h.handle(html_content)
                                f.write(markdown_content)
                            
                            print(f"已保存章节：{current_chapter}")
                        
                        # 开始新的章节
                        current_chapter = heading_text
                        current_content = []
                    
                    # 添加当前标题和内容到章节
                    current_content.append(str(heading))
                    # 添加标题后的内容
                    next_element = heading.next_sibling
                    while next_element and next_element.name not in ['h1', 'h2', 'h3', 'h4', 'h5', 'h6']:
                        if next_element.name:
                            current_content.append(str(next_element))
                        next_element = next_element.next_sibling
        
        # 保存最后一个章节
        if current_chapter:
            safe_filename = re.sub(r'[\\/*?:"<>|]', '_', current_chapter)
            safe_filename = re.sub(r'\s+', '_', safe_filename)
            
            output_path = os.path.join(output_dir, f"{safe_filename}.md")
            with open(output_path, 'w', encoding='utf-8') as f:
                f.write(f"# {current_chapter}\n\n")
                # 使用 html2text 转换内容
                html_content = '\n'.join(current_content)
                # 将 img 标签转换为 Markdown 图片语法
                html_content = convert_img_to_markdown(html_content)
                markdown_content = h.handle(html_content)
                f.write(markdown_content)
            
            print(f"已保存章节：{current_chapter}")
        
        print(f"\n转换完成！所有章节已保存至目录：{output_dir}")
        print(f"图片已保存至目录：{images_dir}")
    
    except Exception as e:
        print(f"转换过程中出现错误：{str(e)}")

if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='将 EPUB 文件转换为 Markdown 格式，每个章节保存为单独的文件')
    parser.add_argument('epub_path', help='EPUB 文件路径')
    parser.add_argument('-o', '--output', help='输出目录路径（可选）')
    
    args = parser.parse_args()
    epub_to_markdown(args.epub_path, args.output) 